package com.alg.local.gateway.filter;

import cn.hutool.json.JSONUtil;
import com.alg.local.gateway.entity.EncrProps;
import com.alg.local.gateway.util.GatewayUtil;
import com.alg.local.gateway.util.SymmetricEncryptionUtilCbc;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.PathMatcher;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;

/**
 * @author al
 * @date 2022/1/21 15:24
 * @description Authentication【认证】、Authorization【授权】
 */
@Component
@Slf4j
public class DecryptRequestBodyFilter implements GlobalFilter, Ordered {

    @Autowired
    private EncrProps encrProps;
    private final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();
    private final PathMatcher pathMatcher = new AntPathMatcher();
    @Autowired
    RedissonClient redissonClient;
    @Autowired
    GatewayUtil gatewayUtil;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        if (encrProps.isEnable()) {
            ServerHttpRequest request = exchange.getRequest();

            boolean canSkip = false;
            List<String> contentType = request.getHeaders().get("Content-Type");
            String str = JSONUtil.toJsonStr(contentType);
            if (contentType != null && str.contains("multipart/form-data")) {
                return chain.filter(exchange);
            }

            log.debug("免校验的URL（配置的白名单+swagger）");
            String requestUrl = request.getPath().toString();

            List<String> ignoreUrlList = encrProps.getIgnoreUrl();
            if (ignoreUrlList != null && ignoreUrlList.size() > 0) {
                canSkip = ignoreUrlList.stream().anyMatch(urlPattern -> pathMatcher.match(urlPattern, requestUrl));
            }

            List<String> ignoreUrlReqList = encrProps.getIgnoreUrlReq();
            if (ignoreUrlReqList != null && ignoreUrlReqList.size() > 0) {
                canSkip = ignoreUrlReqList.stream().anyMatch(urlPattern -> pathMatcher.match(urlPattern, requestUrl));
            }

            if (canSkip) {
                return chain.filter(exchange);
            }

            ServerRequest serverRequest = ServerRequest.create(exchange, this.messageReaders);
            String method = Objects.requireNonNull(request.getMethod()).name();
            if (RequestMethod.POST.name().equals(method)) {
                log.info("[Gateway - RequestBodyDecoder] 请求地址：{}，请求资源路径：{}", request.getRemoteAddress(), request.getPath());
                log.debug("获取加密请求体数据 & 解密请求体数据");
                Mono<String> modifiedBody = serverRequest.bodyToMono(String.class)
                        .doOnNext(s -> log.info(String.format("解密前请求body：%s", s)))
                        .flatMap(originalBody -> decryptBody().apply(exchange, Mono.just(originalBody)))
                        .doOnNext(s -> log.info(String.format("解密后请求body：%s", s)));

                BodyInserter<Mono<String>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
                HttpHeaders headers = new HttpHeaders();
                headers.putAll(exchange.getRequest().getHeaders());
                headers.remove(HttpHeaders.CONTENT_LENGTH);
                CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);

                return bodyInserter.insert(outputMessage, new BodyInserterContext())
                        .then(Mono.defer(() -> {
                            ServerHttpRequest decorator = decorate(exchange, headers, outputMessage);
                            return chain.filter(exchange.mutate().request(decorator).build());
                        }));

            }
        }
        return chain.filter(exchange);
    }

    @Override
    public int getOrder() {
        return -1;
    }

    private BiFunction<ServerWebExchange, Mono<String>, Mono<String>> decryptBody() {
        return (exchange, body) -> {
            AtomicReference<String> result = new AtomicReference<>();
            body.subscribe(value -> {
                        try {
                            ServerHttpRequest request = exchange.getRequest();
                            Map<String, String> aesKey = gatewayUtil.getClientAesKey(request);
                            result.set(SymmetricEncryptionUtilCbc.decrypt(aesKey.get("key"), aesKey.get("iv"), value));
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    },
                    e -> log.error(e.getMessage(), e)
            );
            return Mono.just(result.get());
        };
    }

    private ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers,
                                                CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                if (contentLength > 0) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                }
                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };

    }
}
